import cv2
import numpy as np
import glob, os

# 1. 参数
pattern_size = (8, 5)
square_size  = 25.0

# 2. 世界坐标里的角点
objp = np.zeros((pattern_size[1]*pattern_size[0], 3), np.float32)
objp[:, :2] = np.indices(pattern_size).T.reshape(-1, 2) * square_size

objpoints = []
imgpoints = []

# 3. 读图、找角点
images = glob.glob(os.path.join("calib_images", "*.jpg"))
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    found, corners = cv2.findChessboardCorners(
        gray, pattern_size,
        cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE
    )

    # ←—— 先做 guard，确保后面 cornerSubPix 拿到的数组非空
    if not found or corners is None or len(corners) == 0:
        print(f"⚠️ {os.path.basename(fname)} 未检测到角点，跳过")
        continue

    # 4. 只有真的找到角点，才细化到亚像素
    corners = cv2.cornerSubPix(
        gray, corners, (11,11), (-1,-1),
        criteria=(
            cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
            30, 0.001
        )
    )

    objpoints.append(objp)
    imgpoints.append(corners)

    # 可视化确认
    cv2.drawChessboardCorners(img, pattern_size, corners, found)
    cv2.imshow("Corners", img)
    cv2.waitKey(100)

cv2.destroyAllWindows()
print(f"✅ 成功检测到 {len(objpoints)} 张标定图")

if len(objpoints) == 0:
    print("❌ 一张都没检测到！请检查棋盘格生成和角点检测参数。")
    exit(1)

# 5. 调用calibrateCamera
h, w = gray.shape[:2]
ret, cameraMatrix, distCoeffs, _, _ = cv2.calibrateCamera(
    objpoints, imgpoints, (w, h), None, None
)
print(f"标定完成，重投影误差 {ret:.4f}")
print("cameraMatrix:\n", cameraMatrix)
print("distCoeffs:\n", distCoeffs.ravel())

# 6. 保存
np.savez("camera_calib.npz",
         cameraMatrix=cameraMatrix,
         distCoeffs=distCoeffs,
         imageSize=(w, h))
print("已保存 camera_calib.npz")
